{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "doNFRjPqiBhM",
    "colab_type": "code",
    "outputId": "8e86b76a-fae7-4a0b-db24-447986c06f96",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119.0
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Archive:  uncased_L-12_H-768_A-12.zip\n",
      "  inflating: uncased_L-12_H-768_A-12/bert_model.ckpt.meta  \n",
      "  inflating: uncased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001  \n",
      "  inflating: uncased_L-12_H-768_A-12/vocab.txt  \n",
      "  inflating: uncased_L-12_H-768_A-12/bert_model.ckpt.index  \n",
      "  inflating: uncased_L-12_H-768_A-12/bert_config.json  \n"
     ]
    }
   ],
   "source": [
    "# @title Preparation\n",
    "!pip install -q keras-bert\n",
    "!wget -q https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\n",
    "!unzip -o uncased_L-12_H-768_A-12.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "KUQ8UtquieFj",
    "colab_type": "code",
    "colab": {}
   },
   "outputs": [],
   "source": [
    "# @title Environment\n",
    "import os\n",
    "\n",
    "pretrained_path = 'uncased_L-12_H-768_A-12'\n",
    "config_path = os.path.join(pretrained_path, 'bert_config.json')\n",
    "checkpoint_path = os.path.join(pretrained_path, 'bert_model.ckpt')\n",
    "vocab_path = os.path.join(pretrained_path, 'vocab.txt')\n",
    "\n",
    "# TF_KERAS must be added to environment variables in order to use TPU\n",
    "os.environ['TF_KERAS'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "exDvuSwPevQP",
    "colab_type": "code",
    "colab": {}
   },
   "outputs": [],
   "source": [
    "# @title Initialize TPU Strategy\n",
    "\n",
    "import tensorflow as tf\n",
    "from keras_bert import get_custom_objects\n",
    "\n",
    "TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
    "resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)\n",
    "tf.contrib.distribute.initialize_tpu_system(resolver)\n",
    "strategy = tf.contrib.distribute.TPUStrategy(resolver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "sVTPNxOyj4HJ",
    "colab_type": "code",
    "colab": {}
   },
   "outputs": [],
   "source": [
    "# @title Load Basic Model\n",
    "import codecs\n",
    "from keras_bert import load_trained_model_from_checkpoint\n",
    "\n",
    "token_dict = {}\n",
    "with codecs.open(vocab_path, 'r', 'utf8') as reader:\n",
    "    for line in reader:\n",
    "        token = line.strip()\n",
    "        token_dict[token] = len(token_dict)\n",
    "\n",
    "with strategy.scope():\n",
    "    model = load_trained_model_from_checkpoint(config_path, checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "xioN-O_vtztC",
    "colab_type": "code",
    "outputId": "e147fde2-d1cb-439b-d63f-a649be4a0def",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 377.0
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[CLS]', 'from', 'that', 'day', 'forth', '.', '.', '.', 'my', 'arm', 'changed', '.', '.', '.', 'and', 'a', 'voice', 'echoed', '[SEP]']\n",
      "[CLS] [0.24250675737857819, 0.04605229198932648, -0.24484458565711975, -0.5553151369094849, -0.16091349720954895, -0.046603765338659286, 0.2648216784000397, 0.4632352590560913, -0.15888850390911102, -0.5054463148117065, 0.1872796267271042, -0.0844820961356163, -0.10551590472459793, 1.1445354223251343, 0.35309019684791565, 0.4587448537349701, 0.22255975008010864, 0.23128646612167358, 0.17905741930007935]\n",
      "from [0.2858668565750122, 0.12927496433258057, 0.08937370777130127, -0.06506256759166718, -0.18307062983512878, 0.4357893466949463, -0.4666714668273926, 1.1149680614471436, 0.26170825958251953, -1.0477269887924194, -0.7197380661964417, -0.30874621868133545, 0.3589649498462677, 0.43190720677375793, 0.510287880897522, 0.4445205330848694, 0.6695327162742615, 0.11726009100675583, 0.34817394614219666]\n",
      "that [-0.7514970302581787, 0.14548861980438232, 0.32245880365371704, -0.043174318969249725, 0.5105547904968262, 0.6441463828086853, 0.3476734757423401, 1.8521389961242676, -0.3416602611541748, -0.32603877782821655, 0.3870062232017517, -0.781717836856842, 0.7790629267692566, 0.38004353642463684, -0.026217259466648102, 1.1450964212417603, 0.8676536083221436, -0.006007872521877289, 0.4079648554325104]\n",
      "day [-0.5767358541488647, 0.16873040795326233, 0.7379938364028931, -0.43202054500579834, -0.2541211247444153, -0.12445046007633209, 0.0982133075594902, 0.6879273056983948, -0.5429892539978027, -1.1291855573654175, 0.10011367499828339, -0.7181406021118164, 0.11406952142715454, 0.7394105792045593, -0.14538182318210602, 0.8026691675186157, -0.0491754524409771, -0.5183604955673218, 0.5328009128570557]\n",
      "forth [0.6604690551757812, 0.44872960448265076, 0.22819921374320984, -1.302493691444397, -0.5676481127738953, 0.7631499171257019, 0.31730082631111145, 0.9745714664459229, 0.7384849786758423, -0.12949073314666748, 0.6890379786491394, -0.2063390463590622, 0.7543689608573914, 1.060841679573059, 0.2897346615791321, 1.460550308227539, 0.1994684338569641, -0.3415696620941162, 0.3051512837409973]\n",
      ". [0.7019387483596802, -1.0037682056427002, 0.5318847894668579, -0.2360236793756485, 0.7502506375312805, 0.3656868040561676, 0.47489359974861145, 1.2967716455459595, 0.2107715904712677, -1.135636568069458, -0.2472764253616333, -0.8922532796859741, -0.5542468428611755, 1.2267390489578247, 0.49479609727859497, 0.5181231498718262, 0.246188685297966, -0.028614960610866547, 0.36292847990989685]\n",
      ". [0.6940287351608276, -0.9777721762657166, 0.8698938488960266, 0.10904823988676071, 0.6616824269294739, -0.5494440793991089, 0.6290906667709351, 0.9006298780441284, -0.3196481764316559, -1.679004192352295, -0.45238935947418213, -1.0310773849487305, -0.4837859869003296, 0.8608845472335815, 0.3741556406021118, 0.7373818755149841, 0.06710120290517807, -0.005348920822143555, 0.5634738206863403]\n",
      ". [0.4909120500087738, -0.5309440493583679, 0.4432699978351593, 0.08617030084133148, 0.07127449661493301, 0.002989581786096096, 0.5167111158370972, 0.9511327743530273, 0.20563672482967377, -1.380079984664917, -0.13269132375717163, -0.5576552152633667, -0.6243664026260376, 0.33799853920936584, 0.7936240434646606, 0.006669636815786362, 0.23250898718833923, -0.2893766462802887, 0.013550074771046638]\n",
      "my [0.1981644332408905, 0.18493112921714783, 0.6035155653953552, -0.4148944020271301, -0.26000604033470154, 0.7049614787101746, 0.33354446291923523, 1.4654905796051025, -0.03241407871246338, -0.844508171081543, 0.22912199795246124, -0.4619494676589966, 0.1021852195262909, 0.6551948189735413, 0.4241233468055725, 0.28182658553123474, 0.8504406809806824, -0.7569733262062073, -0.03992771729826927]\n",
      "arm [0.1387634128332138, -0.18610626459121704, 0.043312449008226395, -0.30445870757102966, 0.25981974601745605, 0.43626904487609863, -0.1595546454191208, 0.6088374853134155, -0.2485939860343933, -0.947198748588562, -0.2202623188495636, -0.43322864174842834, 0.35308578610420227, 0.5505412220954895, -0.27582991123199463, 0.8192852735519409, 0.7734602093696594, 0.19990915060043335, -0.1323583722114563]\n",
      "changed [0.6847427487373352, -0.06571120023727417, 0.740653932094574, -0.3307543098926544, 0.5178371667861938, 0.4792756140232086, -0.1748911589384079, 0.7557358145713806, -0.8578546643257141, 0.029164336621761322, 0.8551721572875977, -0.9205213189125061, 0.11773020774126053, 0.9219750165939331, 0.4660824239253998, 0.33519038558006287, 0.5005586743354797, -0.4561012089252472, -0.8589122891426086]\n",
      ". [0.691735029220581, -0.9538367390632629, 0.6850607991218567, -0.33083242177963257, 0.6519606113433838, 0.18196794390678406, 0.14063923060894012, 1.0543153285980225, 0.09589402377605438, -1.1536606550216675, -0.3813367486000061, -0.7072638869285583, -0.5489786863327026, 1.2570656538009644, 0.4382341802120209, 0.4493916630744934, 0.46416541934013367, -0.1259787380695343, 0.13374735414981842]\n",
      ". [0.49735528230667114, -0.824055016040802, 1.0104886293411255, 0.1664600521326065, 0.6780160069465637, -0.6264775991439819, 0.3941092789173126, 0.7728183269500732, -0.33566147089004517, -1.6417758464813232, -0.3855264186859131, -1.174580454826355, -0.48833924531936646, 0.8608969449996948, 0.4212232530117035, 0.6263501048088074, 0.11568772792816162, -0.18229623138904572, 0.4028193950653076]\n",
      ". [0.5248441100120544, -0.8752838373184204, 0.8074136972427368, -0.30946090817451477, 0.5306240916252136, -0.11764449626207352, 0.6163278222084045, 0.8180994987487793, 0.0724523738026619, -0.7624974250793457, -0.3387123942375183, -0.6998348832130432, -0.7396746873855591, 0.9056875705718994, 0.3457651436328888, 0.07841179519891739, -0.011172758415341377, -0.5140762329101562, 0.0032963640987873077]\n",
      "and [-0.1338833123445511, -0.19814497232437134, -0.16618409752845764, 0.07875220477581024, -0.1724027693271637, 0.7233635783195496, -0.1553294062614441, 1.6001187562942505, 0.045651875436306, -1.3001012802124023, 0.1902690976858139, -0.34323716163635254, 0.19466280937194824, 0.6093133687973022, 0.4001396894454956, 0.5367133617401123, 0.609708845615387, -0.5849275588989258, -0.5499515533447266]\n",
      "a [-0.07800914347171783, 0.08851209282875061, -0.12622755765914917, -0.36778488755226135, 0.5426658391952515, 0.3223933279514313, 0.3393816351890564, 0.8235282897949219, -0.43569034337997437, -0.829987645149231, 0.05009227991104126, -0.7572609186172485, -0.5095729827880859, 1.2285490036010742, -0.5328006148338318, -0.08091999590396881, 0.8362939357757568, -0.3866221606731415, -0.07156673073768616]\n",
      "voice [0.43880149722099304, -0.12388401478528976, -0.09448350965976715, -0.4879209101200104, 0.6075019836425781, 0.32898378372192383, 0.6265825033187866, 0.728482723236084, -0.46667003631591797, -0.5817348957061768, -0.4383503198623657, -0.31208136677742004, -0.4085504412651062, 0.6028963327407837, -0.6348038911819458, -0.43311747908592224, -0.12351862341165543, -0.07814896106719971, -0.16226796805858612]\n",
      "echoed [0.6161936521530151, -0.22712212800979614, -0.0740918517112732, -0.7913801074028015, 0.09446655958890915, 0.6521013975143433, -0.08286319673061371, 0.4983268082141876, -0.629666268825531, -0.40694940090179443, 0.12792044878005981, -0.11819512397050858, -0.4466920495033264, 0.8771257996559143, -0.8426066637039185, -0.34973663091659546, 0.11151856929063797, -0.30014142394065857, -0.5564588308334351]\n",
      "[SEP] [0.7558457851409912, 0.36590367555618286, -0.17892876267433167, 0.39153221249580383, -0.5322665572166443, -0.613459587097168, 0.4913185238838196, -0.35385745763778687, 0.4216068983078003, -0.1137060597538948, 0.17552773654460907, -0.10072151571512222, 0.1765543669462204, 0.14826729893684387, -0.25414443016052246, -0.06873864680528641, 0.23133708536624908, -0.04865124076604843, 0.37981322407722473]\n"
     ]
    }
   ],
   "source": [
    "# @title Extraction\n",
    "import numpy as np\n",
    "from keras_bert import Tokenizer\n",
    "\n",
    "tokenizer = Tokenizer(token_dict)\n",
    "text = 'From that day forth... my arm changed... and a voice echoed'\n",
    "tokens = tokenizer.tokenize(text)\n",
    "indices, segments = tokenizer.encode(first=text, max_len=512)\n",
    "print(tokens)\n",
    "\n",
    "predicts = model.predict([np.array([indices] * 8), np.array([segments] * 8)])[0]\n",
    "    \n",
    "for i, token in enumerate(tokens):\n",
    "    print(token, predicts[i].tolist()[:19])"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "keras_bert_load_and_extract_tpu.ipynb",
   "version": "0.3.2",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "TPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
